package com.fr.football;

import java.util.HashMap;
import java.util.Map;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.mllib.evaluation.MulticlassMetrics;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.tree.RandomForest;
import org.apache.spark.mllib.tree.model.DecisionTreeModel;
import org.apache.spark.mllib.tree.model.RandomForestModel;
import org.apache.spark.util.Utils;

import scala.Tuple2;

public class Train {

	public static void main(String[] args) {
		SparkConf conf = new SparkConf().setMaster("local").setAppName("fr");
		JavaSparkContext jsc = new JavaSparkContext(conf);
		JavaRDD<LabeledPoint> data = readData(jsc);
		// 将数据分割为训练集、交叉检验集(CV)和测试集

		JavaRDD<LabeledPoint>[] splitArray = data.randomSplit(new double[] { 0.8, 0.1, 0.1 });

		JavaRDD<LabeledPoint> trainData = splitArray[0];

		trainData.cache();

		JavaRDD<LabeledPoint> cvData = splitArray[1];

		cvData.cache();

		JavaRDD<LabeledPoint> testData = splitArray[2];

		testData.cache();
		
		// 构建DecisionTreeModel

//		DecisionTreeModel model = DecisionTree.trainClassifier(data, 7, new HashMap<Integer, Integer>(), "gini", 4, 100);
//		MulticlassMetrics metrics = getMetrics(model, cvData);
//		System.out.println(metrics.confusionMatrix());
//		System.out.println(metrics.precision());
		Map<Integer, Integer> map = new HashMap<Integer, Integer>();
		map.put(10, 4);
		map.put(11, 40);
		//构建RandomForestModel
		RandomForestModel model = RandomForest.trainClassifier(trainData, 3, map, 20, "auto", "entropy", 30, 300, Utils.random().nextInt());

		//用CV集来计算结果模型的指标
		MulticlassMetrics metrics = getMetrics(model, cvData);
		System.out.println(metrics.precision());
		// 用CV集来计算结果模型的指标
		
	}

	public static JavaRDD<LabeledPoint> readData(JavaSparkContext jsc) {

		JavaRDD<String> rawData = jsc.textFile("F:\\data\\1.txt");

		JavaRDD<LabeledPoint> data = rawData.map(line -> {

			String[] values = line.split("\t");

			double[] features = new double[values.length - 1];

			for (int i = 0; i < values.length - 1; i++) {
				double value = Double.parseDouble(values[i]);
				if (value == 0.00d) {
					continue;
				}
				features[i] = value;
			}
			Vector featureVector = Vectors.dense(features);
			Double label = (double) (Double.parseDouble(values[values.length - 1]));
			return new LabeledPoint(label, featureVector);
		});
		return data;
	}

	public static MulticlassMetrics getMetrics(DecisionTreeModel model, JavaRDD<LabeledPoint> data) {
		JavaPairRDD<Object, Object> predictionsAndLabels = data.mapToPair(example -> {
			return new Tuple2<Object, Object>(model.predict(example.features()), example.label());
		});
		return new MulticlassMetrics(JavaPairRDD.toRDD(predictionsAndLabels));
	}
	
	public static MulticlassMetrics getMetrics(RandomForestModel model, JavaRDD<LabeledPoint> data) {
		JavaPairRDD<Object, Object> predictionsAndLabels = data.mapToPair(example -> {
			return new Tuple2<Object, Object>(model.predict(example.features()), example.label());
		});
		return new MulticlassMetrics(JavaPairRDD.toRDD(predictionsAndLabels));
	}

}
